import uuid
from contextvars import ContextVar
from quart import request, g
from typing import Optional

# 使用ContextVar来存储请求ID，确保线程安全
request_id_var: ContextVar[Optional[str]] = ContextVar('request_id', default=None)


class RequestIdMiddleware:
    """
    请求ID中间件
    为每个请求生成唯一的ID，并存储在上下文中
    """
    
    def __init__(self, app):
        self.app = app
        self.init_app(app)
    
    def init_app(self, app):
        """初始化中间件"""
        app.before_request(self.before_request)
        app.after_request(self.after_request)
    
    async def before_request(self):
        """请求前处理：生成或获取请求ID"""
        # 尝试从请求头获取请求ID
        request_id = request.headers.get('X-Request-ID')
        
        # 如果没有提供请求ID，则生成一个新的
        if not request_id:
            request_id = str(uuid.uuid4())
        
        # 存储到上下文变量中
        request_id_var.set(request_id)
        
        # 也存储到g对象中，方便在视图函数中使用
        g.request_id = request_id
    
    async def after_request(self, response):
        """请求后处理：在响应头中添加请求ID"""
        request_id = get_request_id()
        if request_id:
            response.headers['X-Request-ID'] = request_id
        return response


def get_request_id() -> Optional[str]:
    """
    获取当前请求的ID
    """
    return request_id_var.get()


def generate_trace_id() -> str:
    """
    生成追踪ID（与请求ID保持一致）
    """
    request_id = get_request_id()
    if request_id:
        return request_id
    return str(uuid.uuid4())